﻿Imports System.Runtime.CompilerServices

Public Module TreeViewExtensions
    <Extension()> _
     Public Function ContainerFromItem(ByVal ThisTreeView As TreeView, ByVal item As Object) As TreeViewItem
        Dim containerThatMightContainItem As TreeViewItem = ThisTreeView.ItemContainerGenerator.ContainerFromItem(item)
        If containerThatMightContainItem IsNot Nothing Then
            Return containerThatMightContainItem
        Else
            Return ContainerFromItem(ThisTreeView.ItemContainerGenerator, ThisTreeView.Items, item)
        End If
    End Function
    Private Function ContainerFromItem(ByVal parentItemContainerGenerator As ItemContainerGenerator, ByVal ThisItemCollection As ItemCollection, ByVal item As Object) As TreeViewItem
        For Each curChildItem As Object In ThisItemCollection
            Dim parentContainer As TreeViewItem = parentItemContainerGenerator.ContainerFromItem(curChildItem)
            If parentContainer Is Nothing Then
                Return Nothing
            End If
            Dim containerThatMightContainItem As TreeViewItem = parentContainer.ItemContainerGenerator.ContainerFromItem(item)
            If containerThatMightContainItem IsNot Nothing Then
                Return containerThatMightContainItem
            End If
            Dim recursionResult As TreeViewItem = ContainerFromItem(parentContainer.ItemContainerGenerator, parentContainer.Items, item)
            If recursionResult IsNot Nothing Then
                Return recursionResult
            End If
        Next
        Return Nothing
    End Function
    <Extension()> _
    Public Function ItemFromContainer(ByVal ThisTreeView As TreeView, ByVal container As TreeViewItem) As Object
        Dim itemThatMightBelongToContainer As TreeViewItem = ThisTreeView.ItemContainerGenerator.ItemFromContainer(container)
        If itemThatMightBelongToContainer IsNot Nothing Then
            Return itemThatMightBelongToContainer
        Else
            Return ItemFromContainer(ThisTreeView.ItemContainerGenerator, ThisTreeView.Items, container)
        End If
    End Function
    Private Function ItemFromContainer(ByVal parentItemContainerGenerator As ItemContainerGenerator, ByVal ThisItemCollection As ItemCollection, ByVal container As TreeViewItem) As Object
        For Each curChildItem As Object In ThisItemCollection
            Dim parentContainer As TreeViewItem = parentItemContainerGenerator.ContainerFromItem(curChildItem)
            If parentContainer Is Nothing Then
                Return Nothing
            End If
            Dim itemThatMightBelongToContainer As TreeViewItem = parentContainer.ItemContainerGenerator.ItemFromContainer(container)
            If itemThatMightBelongToContainer IsNot Nothing Then
                Return itemThatMightBelongToContainer
            End If
            Dim recursionResult As TreeViewItem = ItemFromContainer(parentContainer.ItemContainerGenerator, parentContainer.Items, container)
            If recursionResult IsNot Nothing Then
                Return recursionResult
            End If
        Next
        Return Nothing
    End Function
    <Extension()> _
    Public Function FindContainer(ByVal thisTreeView As TreeView, ByVal condition As Predicate(Of TreeViewItem)) As TreeViewItem
        Return FindContainer(thisTreeView.ItemContainerGenerator, thisTreeView.Items, condition)
    End Function
    Private Function FindContainer(ByVal parentItemContainerGenerator As ItemContainerGenerator, ByVal ThisItemCollection As ItemCollection, ByVal condition As Predicate(Of TreeViewItem)) As TreeViewItem
        For Each curChildItem As Object In ThisItemCollection
            Dim containerThatMightMeetTheCondition As TreeViewItem = parentItemContainerGenerator.ContainerFromItem(curChildItem)
            If containerThatMightMeetTheCondition Is Nothing Then
                Return Nothing
            End If
            If condition(containerThatMightMeetTheCondition) Then
                Return containerThatMightMeetTheCondition
            End If
            Dim recursionResult As TreeViewItem = FindContainer(containerThatMightMeetTheCondition.ItemContainerGenerator, containerThatMightMeetTheCondition.Items, condition)
            If recursionResult IsNot Nothing Then
                Return recursionResult
            End If
        Next
        Return Nothing
    End Function
End Module